import numpy as np
from scipy.optimize import linear_sum_assignment


def iou(bb_test, bb_gt):
    """
    Computes IUO between two bboxes in the form [x1,y1,x2,y2]
    """
    xx1 = np.maximum(bb_test[0], bb_gt[0])
    yy1 = np.maximum(bb_test[1], bb_gt[1])
    xx2 = np.minimum(bb_test[2], bb_gt[2])
    yy2 = np.minimum(bb_test[3], bb_gt[3])
    w = np.maximum(0., xx2 - xx1)
    h = np.maximum(0., yy2 - yy1)
    wh = w * h
    r = wh / ((bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1]) +
              (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1]) - wh)
    return r


def associate_detections_to_trackers(detections, trackers, iou_threshold):
    """
    Assigns detections to tracked object (both represented as bounding boxes)

    Returns 3 lists of matches, unmatched_detections and unmatched_trackers
    """
    if len(trackers) == 0:
        matches = np.empty((0, 2), dtype=np.int32)
        unmatched_detections = np.arange(len(detections))
        unmatched_trackers = np.empty((0, 5), dtype=np.int32)
        return matches, unmatched_detections, unmatched_trackers

    iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)

    for d, det in enumerate(detections):
        for t, trk in enumerate(trackers):
            iou_matrix[d, t] = iou(det, trk)
    """
    The linear assignment module tries to minimise the total assignment cost.
    In our case we pass -iou_matrix as we want to maximise the total IOU 
        between track predictions and the frame detection.
    """
    # matched_indices is a tuple, corresponding to row_ind, col_ind
    matched_indices = linear_sum_assignment(-iou_matrix)

    unmatched_detections = []
    for d, det in enumerate(detections):
        # if d not in matched_indices[:, 0]:
        if d not in matched_indices[0]:
            unmatched_detections.append(d)
    unmatched_trackers = []
    for t, trk in enumerate(trackers):
        # if t not in matched_indices[:, 1]:
        if t not in matched_indices[1]:
            unmatched_trackers.append(t)

    # filter out matched with low IOU
    matches = []
    for y, x in zip(matched_indices[0], matched_indices[1]):
        if iou_matrix[y, x] < iou_threshold:
            unmatched_detections.append(y)
            unmatched_trackers.append(x)
        else:
            matches.append(np.array([[y, x]]))
    if len(matches) == 0:
        matches = np.empty((0, 2), dtype=np.int32)
    else:
        matches = np.concatenate(matches, axis=0)

    unmatched_detections = np.array(unmatched_detections)
    unmatched_trackers = np.array(unmatched_trackers)
    return matches, unmatched_detections, unmatched_trackers
